#include <ros/ros.h>

#include <sensor_msgs/point_cloud_conversion.h>
#include <sensor_msgs/PointCloud2.h>
#include "sensor_msgs/Imu.h"
#include <visualization_msgs/Marker.h>

// PCL specific includes
#include <pcl/io/pcd_io.h>
#include <pcl_conversions/pcl_conversions.h>
#include <pcl/point_cloud.h>
#include <pcl/point_types.h>
#include <pcl/common/transforms.h>
#include <pcl_ros/transforms.h>
#include "pcl_ros/impl/transforms.hpp"

#include <iostream>
#include <chrono>
#include <vector>
#include <math.h>

#include <geometry_msgs/TwistStamped.h>
#include <tf/transform_broadcaster.h>
#include <tf/transform_listener.h>
#include "tf/transform_datatypes.h"
#include <tf2/LinearMath/Quaternion.h>
#include <tf2_ros/transform_broadcaster.h>
#include <object_tracking/trackbox.h>
#include <object_tracking/resbox.h>
#include <object_tracking/resboxes.h>

#include "ImmUkf.hpp"
#include "vis_bbox.h"

using namespace std;
using namespace Eigen;
using namespace pcl;


// private parameters
string lidar_name = "velo_link";
string global_name = "global";
string params_path;

// global listener: listening transformation from lidar to global 
tf::TransformListener* tran;

// pub result
ros::Publisher res_pub;
//ros::Publisher vis_pub;

// tracking object !!
ImmUkf Tracking;

using time_point = std::chrono::time_point<std::chrono::high_resolution_clock>;
time_point st;
time_point ed;
auto& CLOCK = std::chrono::high_resolution_clock::now;
// print time
void logTime(string info)
{
	std::chrono::duration<double, std::milli> fp_ms = ed - st;
	string str = info + " cost %f ms";
	ROS_DEBUG_NAMED("main", str.c_str(), fp_ms.count());
}

// rearrange the input
vector<PointCloud<PointXYZ>> formatInput(const object_tracking::trackbox& input)
{
	
	vector<PointCloud<PointXYZ>> bBoxes;
	int box_num = input.box_num;
	bBoxes.reserve(box_num);

	for(int box_i = 0;box_i < box_num; box_i++)
	{
		PointCloud<PointXYZ> oneBbox;
		PointXYZ o;
		o.x = input.x1[3*box_i];
		o.y = input.x1[3*box_i + 1];
		o.z = input.x1[3*box_i + 2];
		oneBbox.push_back(o);
		o.x = input.x2[3*box_i];
		o.y = input.x2[3*box_i + 1];
		o.z = input.x2[3*box_i + 2];   
		oneBbox.push_back(o);
		o.x = input.x3[3*box_i + 0];
		o.y = input.x3[3*box_i + 1];
		o.z = input.x3[3*box_i + 2];   
		oneBbox.push_back(o);
		o.x = input.x4[3*box_i + 0];
		o.y = input.x4[3*box_i + 1];
		o.z = input.x4[3*box_i + 2];   
		oneBbox.push_back(o);
		o.x = input.y1[3*box_i + 0];
		o.y = input.y1[3*box_i + 1];
		o.z = input.y1[3*box_i + 2];   
		oneBbox.push_back(o);
		o.x = input.y2[3*box_i + 0];
		o.y = input.y2[3*box_i + 1];
		o.z = input.y2[3*box_i + 2];   
		oneBbox.push_back(o);
		o.x = input.y3[3*box_i + 0];
		o.y = input.y3[3*box_i + 1];
		o.z = input.y3[3*box_i + 2];   
		oneBbox.push_back(o);
		o.x = input.y4[3*box_i + 0];
		o.y = input.y4[3*box_i + 1];
		o.z = input.y4[3*box_i + 2];   
		oneBbox.push_back(o);
		bBoxes.emplace_back(oneBbox);
	}

	return bBoxes;
}

void pubResBoxes(ros::Time stamp, string frame_id, const TrackingResults& tres, const vector<PointCloud<PointXYZ>>& trackingBBoxes, const vector<vector<Vector3d>>& predTraj)
{
	object_tracking::resboxes rbs;
	rbs.header.stamp = stamp;
	rbs.header.frame_id = frame_id;
	
	int n = trackingBBoxes.size();
	rbs.rb.reserve(n);	
	for(int i=0; i<n; i++){
		// push id and color
		object_tracking::resbox erb;
		erb.id = tres.id[i];
		erb.vel = tres.vel[i];
		erb.yaw = tres.bxyzAndYaw[i](3);

		for(int j=0; j<3; j++) 	erb.xyz.emplace_back(tres.bxyzAndYaw[i](j));
		for(int j=0; j<3; j++) 	erb.color.emplace_back(tres.color[i](j));

		const auto& e = trackingBBoxes[i];
		for(int j=0; j<8; j++){
			erb.bx.emplace_back(e[j].x);
			erb.by.emplace_back(e[j].y);
			erb.bz.emplace_back(e[j].z);
		}
		const auto& e1 = predTraj[i];
		// 没办法，这里暂时先这样
		//if(e1.size() < 2){
			//ROS_WARN("traj empty bug not fixed!!!");	
			//return;
		//}
		for(int j=0; j<e1.size(); j++){
			erb.traj_x.emplace_back(e1[j](0));
			erb.traj_y.emplace_back(e1[j](1));
			erb.traj_z.emplace_back(e1[j](2));
		}
		rbs.rb.emplace_back(erb);
	}

	object_tracking::resbox ego;
	for(int i=0; i<tres.egoTraj.size(); i++){
		ego.traj_x.emplace_back(tres.egoTraj[i].x);
		ego.traj_y.emplace_back(tres.egoTraj[i].y);
		ego.traj_z.emplace_back(tres.egoTraj[i].z);
	}
	rbs.ego = std::move(ego);
	res_pub.publish(rbs);
}


// 回调函数-- 话题 track_box
void cloud_cb (const object_tracking::trackbox& input)
{
//-------------------start converting----------ego compensation-------------------------
	ros::Time input_time = input.header.stamp;
	double timestamp = input_time.toSec();  
	ROS_DEBUG_NAMED("main", "get in cloud_cb at time(%f)", timestamp);

	//ros::Time input_time = ros::Time::now();
	//double timestamp = input_time.toSec();  
	// ----------------convert local to global-------------------------
	tran->waitForTransform("/" + global_name, "/" + lidar_name, input_time, ros::Duration(0.1));
	tf::StampedTransform ego_tf;
	try{
		tran->lookupTransform("/" + global_name, "/" + lidar_name, input_time, ego_tf);
	}catch(exception& e) 
	{
		return;
	}

	Vector3d ego;
	tf::Vector3 tmp_ego = ego_tf.getOrigin();
	double r, p, ego_yaw;
	tf::Matrix3x3(ego_tf.getRotation()).getRPY(r, p, ego_yaw);

	ego << tmp_ego.getX(), tmp_ego.getY(), ego_yaw;

	vector<PointCloud<PointXYZ>> bBoxes = formatInput(input);
	
	std_msgs::Header bh;
	bh.frame_id = lidar_name;
	bh.stamp = input_time;
	PCLHeader pcl_bh = pcl_conversions::toPCL(bh);
	//for(int i = 0; i < bBoxes.size(); i++ ){
		//PointCloud<PointXYZ> newBox;
		//bBoxes[i].header = pcl_bh;
		//if(!pcl_ros::transformPointCloud("/" + global_name ,bBoxes[i], newBox, *tran)) 	return;
		//bBoxes[i] = newBox;
	//}
	//-------------------end converting----------ego compensation-------------------------
	st = CLOCK();
	Tracking.setInput(timestamp, ego, bBoxes);
	Tracking.run(timestamp);
	ed = CLOCK();
	logTime("Tracking.run");
	
	// TODO: arrange the result and pub
	// 1. visualize bbox
	// 2. predicted trajectory
	// 3. local cost map
	
	TrackingResults tRes = Tracking.getArrangeResults();

	// turn pc to local frame
	vector<PointCloud<PointXYZ>> localPc;
	for(int i=0; i<tRes.bPtsAndTrajPts.size(); i++)
	{
		PointCloud<PointXYZ> newBox;
		tRes.bPtsAndTrajPts[i].header = pcl_bh;
		tRes.bPtsAndTrajPts[i].header.frame_id = global_name;
		pcl_ros::transformPointCloud("/" + lidar_name, tRes.bPtsAndTrajPts[i], newBox, *tran);
		localPc.emplace_back(newBox);
	}

	// TODO: visualization
	// 1. pub bbox
	// 2. pub trajectory
	
	// retrive bbox and trajectory
	vector<PointCloud<PointXYZ>> trackingBBoxes;
	vector<vector<Vector3d>> predTraj;
	trackingBBoxes.reserve(localPc.size());
	predTraj.reserve(localPc.size());
	for(auto&& e : localPc){
		PointCloud<PointXYZ> tmp;
		tmp.reserve(8);
		tmp.insert(tmp.end(), e.begin(), e.begin()+8);
		trackingBBoxes.emplace_back(tmp);
		
		vector<Vector3d> oneBoxTraj;
		oneBoxTraj.reserve(std::distance(e.begin()+8, e.end()));
		for(auto it=e.begin()+8; it!=e.end(); it++)
		{
			oneBoxTraj.emplace_back(Vector3d(it->x, it->y, it->z));
		}
		predTraj.emplace_back(oneBoxTraj);
	}

	pubResBoxes(input_time, lidar_name, tRes, trackingBBoxes, predTraj);
	
}



int main (int argc, char** argv){
	// Initialize ROS
	ros::init (argc, argv, "obj_track");
	ros::NodeHandle nh;

	// get params
	nh.getParam("lidar_name", lidar_name);
	nh.getParam("global_name", global_name);
	nh.getParam("params_path", params_path);
	
	cv::FileStorage fs(params_path, cv::FileStorage::READ);
	assert(fs.isOpened());

	Tracking = ImmUkf();
	fs["TrackParams"] >> Tracking;
	fs.release();
	
	Tracking.initEgoTracking();

	// listen to global tf tree
	tf::TransformListener lr(ros::Duration(100.0));
	tran=&lr;

	ros::Subscriber sub = nh.subscribe ("track_box", 1, cloud_cb);   //订阅者  track_box -- 话题topic名
	
	res_pub = nh.advertise<object_tracking::resboxes>( "resbox", 1);
	//vis_pub = nh.advertise<visualization_msgs::Marker>( "tracking_markers", 200);

	// Spin
	ros::spin ();
}
